""" Standard incomplete market model """

using LinearAlgebra
include("het.jl") # from het_block import het
include("interpolate.jl") # from utilities import interpolate
include("misc.jl") # from utilities import misc
include("grids.jl") # from utilities import grids

""" Core HetBlock """

function hh_init(a_grid, y, r, eis)
    coh = (1 + r) * reshape(a_grid, 1, size(a_grid)...) .+ reshape(y, 1, size(y)...)
    Va = (1 + r) * (0.1 * coh) .^ (-1 / eis)
    return Va
end

function _hh(Va_p, a_grid, y, r, beta, eis)
    uc_nextgrid = β * Va_p
    c_nextgrid = uc_nextgrid .^ (-eis)
    coh = (1 + r) * reshape(a_grid, 1, size(a_grid)...) .+ reshape(y, 1, size(y)...)
    a = interpolate_y(c_nextgrid .+ a_grid, coh, a_grid)
    setmin(a, a_grid[1])
    c = coh .- a
    Va = (1 + r) * c .^ (-1 / eis)
    return Va, a, v
end

hh = het(_hh; exogenous = "Pi", policy = "a", backward = "Va", backward_init = hh_init)

""" Extended HetBlock with grid and income process inputs added, and example calibration """

function make_grids(rho_e, sd_e, n_e, min_a, max_a, n_a)
    e_grid, _, Pi = markov_rouwenhorst(rho_e, sd_e, n_e)
    a_grid = grids.asset_grid(min_a, max_a, n_a)
    return e_grid, Pi, a_grid
end

function income(w, e_grid)
    y = w .* e_grid
    return y
end

# extend hh with grid and income process inputs added
hh_extended = hh.add_hetinputs([income, make_grids])

function example_calibration()
    return Dict(min_a => 0,
                max_a => 1000,
                rho_e => 0.975,
                sd_e => 0.7,
                n_a => 200,
                n_e => 7,
                w => 1,
                r => 0.01 / 4,
                β => 1 - 0.08 / 4,
                eis => 1
                )
end
